#ifndef _HOOK_H_
#define _HOOK_H_

#include <Windows.h>

#include <map>
#include <memory>
#include <string>
#include <string_view>
#include <vector>

namespace windows
{

class hook
{
public:
    using function_address = PROC;

    struct module_import
    {
        std::string_view module_name;
        std::vector<std::pair<std::string_view, function_address>> replaceable_functions;
    };

    struct import_address
    {
        function_address original{};
        function_address modified{};
    };

    ~hook();

    static std::unique_ptr<hook>& get_or_create(std::vector<module_import> import_functions = {});
    static import_address get_function_address(const std::string module_name, const std::string_view function_name);

private:
    enum class operation_type : uint8_t
    {
        modify,
        cancel
    };

    static constexpr auto kernel_module = std::string_view("Kernel32.dll");

    using function_table = std::map<std::string, import_address, std::less<>>;
    using module_table = std::map<std::string, function_table, std::less<>>;

    explicit hook(const std::vector<module_import>& import_functions);

    void modify_addresses_in_import_tables(const HMODULE module_address, const operation_type operation) noexcept;
    void modify_addresses_in_table(const operation_type operation, PIMAGE_THUNK_DATA iat, const function_table& addresses) noexcept;

    static void process_newly_module(const HMODULE module_address, const DWORD flags = 0) noexcept;
    static import_address get_function_address(const HMODULE module_address, PCSTR function_name);

    template <typename String>
    static HMODULE load_library(String module_path);

    template <typename String>
    static HMODULE load_library_ex(String module_path, const HANDLE file, const DWORD flags);

    static function_address get_procedure_address(const HMODULE module_address, PCSTR function_name);

    module_table replaceable_tables;
    std::map<HMODULE, std::string> module_names;
    const HMODULE current_module;
};

template <typename String>
HMODULE hook::load_library(String module_path)
{
    HMODULE module_address{};

    if constexpr (std::is_same_v<String, PCSTR>)
    {
        const auto original_load_library = get_function_address(::GetModuleHandleA(kernel_module.c_str()), "LoadLibraryA").original;
        module_address = reinterpret_cast<HMODULE (*)(PCSTR)>(original_load_library)(module_path);
    }
    else if constexpr (std::is_same_v<String, PCWSTR>)
    {
        const auto original_load_library = get_function_address(::GetModuleHandleA(kernel_module.c_str()), "LoadLibraryW").original;
        module_address = reinterpret_cast<HMODULE (*)(PCWSTR)>(original_load_library)(module_path);
    }

    hook::process_newly_module(module_address);

    return module_address;
}

template <typename String>
HMODULE hook::load_library_ex(String module_path, const HANDLE file, const DWORD flags)
{
    HMODULE module_address{};

    if constexpr (std::is_same_v<String, PCSTR>)
    {
        const auto original_load_library = get_function_address(::GetModuleHandleA(kernel_module.c_str()), "LoadLibraryExA").original;
        module_address = reinterpret_cast<HMODULE (*)(PCSTR, HANDLE, DWORD)>(original_load_library)(module_path, file, flags);
    }
    else if constexpr (std::is_same_v<String, PCWSTR>)
    {
        const auto original_load_library = get_function_address(::GetModuleHandleA(kernel_module.c_str()), "LoadLibraryExW").original;
        module_address = reinterpret_cast<HMODULE (*)(PCWSTR, HANDLE, DWORD)>(original_load_library)(module_path, file, flags);
    }

    hook::process_newly_module(module_address, flags);

    return module_address;
}

} // namespace windows

#endif
